package main

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
)

const (
	KB = 1024        // 1KB = 1024 bytes
	MB = 1024 * 1024 // 1MB = 1024 KB
)

// 格式化文件大小，自动转换为 KB 或 MB
func formatSize(sizeBytes int64) string {
	sizeKB := sizeBytes / KB
	if sizeKB < KB {
		return fmt.Sprintf("%d KB", sizeKB) // 小于 1024 KB，直接显示 KB
	} else {
		sizeMB := float64(sizeBytes) / float64(MB)
		return fmt.Sprintf("%.2f MB", sizeMB) // 大于等于 1024 KB，转换为 MB 并保留两位小数
	}
}

// 文件或文件夹信息
type fileInfo struct {
	name  string
	size  int64
	isDir bool
}

// 递归遍历文件夹并收集信息
func walkDir(path string, wg *sync.WaitGroup, ch chan<- fileInfo) {
	defer wg.Done()

	entries, err := os.ReadDir(path)
	if err != nil {
		fmt.Printf("Error reading directory: %s\n", err)
		return
	}

	for _, entry := range entries {
		fullPath := filepath.Join(path, entry.Name())
		info, err := entry.Info()
		if err != nil {
			fmt.Printf("Error getting file info: %s\n", err)
			continue
		}

		ch <- fileInfo{name: fullPath, size: info.Size(), isDir: entry.IsDir()}

		if entry.IsDir() {
			wg.Add(1)
			go walkDir(fullPath, wg, ch)
		}
	}
}

// 打印文件或文件夹信息
func printFileOrDirInfo(name string, size int64, prefix string, maxNameLen int) {
	if size < 0 {
		fmt.Printf("%sdir - %-*s - %s\n", prefix, maxNameLen, name, "Calculating...")
	} else {
		fmt.Printf("%sfile - %-*s - %s\n", prefix, maxNameLen, name, formatSize(size))
	}
}

// 主函数
func main() {
	// 默认路径为当前文件夹，默认递归展示所有层级
	path := "."
	maxDepth := -1 // -1 表示递归展示所有层级

	// 解析命令行参数
	if len(os.Args) > 1 {
		for _, arg := range os.Args[1:] {
			if strings.HasPrefix(arg, "-") {
				depth, err := strconv.Atoi(arg[1:])
				if err != nil {
					fmt.Printf("Invalid depth parameter: %s\n", arg)
					return
				}
				maxDepth = depth
			} else {
				path = arg
			}
		}
	}

	// 如果路径是 "."，则获取当前工作目录
	if path == "." {
		wd, err := os.Getwd()
		if err != nil {
			fmt.Printf("Error getting current working directory: %s\n", err)
			return
		}
		path = wd
	}

	// 使用 Channel 和 WaitGroup 并行遍历文件夹
	ch := make(chan fileInfo)
	var wg sync.WaitGroup
	wg.Add(1)
	go walkDir(path, &wg, ch)

	// 收集所有文件信息
	var files []fileInfo
	go func() {
		wg.Wait()
		close(ch)
	}()

	for info := range ch {
		files = append(files, info)
	}

	// 计算最大文件名长度
	maxNameLen := 0
	for _, info := range files {
		name := filepath.Base(info.name)
		if len(name) > maxNameLen {
			maxNameLen = len(name)
		}
	}

	// 构建文件夹大小映射
	sizeMap := make(map[string]int64)
	for _, info := range files {
		if info.isDir {
			continue
		}
		// 将文件大小累加到其所有父文件夹的大小中
		dir := filepath.Dir(info.name)
		for {
			sizeMap[dir] += info.size
			parentDir := filepath.Dir(dir)
			if parentDir == dir {
				break
			}
			dir = parentDir
		}
	}

	// 打印文件树
	printTree(path, "", maxNameLen, sizeMap, files, maxDepth, 0)

	// 输出分隔符
	fmt.Println("\n" + strings.Repeat("=", 50))
	fmt.Println("Largest 5 files:")

	// 筛选出文件并排序
	var fileList []fileInfo
	for _, info := range files {
		if !info.isDir {
			fileList = append(fileList, info)
		}
	}

	// 按文件大小排序
	sort.Slice(fileList, func(i, j int) bool {
		return fileList[i].size > fileList[j].size
	})

	// 输出最大的 5 个文件
	for i := 0; i < len(fileList) && i < 5; i++ {
		file := fileList[i]
		fmt.Printf("%s - %s\n", file.name, formatSize(file.size))
	}
}

// 递归打印文件夹树
func printTree(path string, prefix string, maxNameLen int, sizeMap map[string]int64, files []fileInfo, maxDepth int, currentDepth int) {
	fileInfo, err := os.Stat(path)
	if err != nil {
		fmt.Printf("Error getting file info: %s\n", err)
		return
	}

	name := filepath.Base(path)
	if fileInfo.IsDir() {
		size := sizeMap[path]
		fmt.Printf("%sdir - %-*s - %s\n", prefix, maxNameLen, name, formatSize(size))
		if maxDepth == -1 || currentDepth < maxDepth {
			entries, err := os.ReadDir(path)
			if err != nil {
				fmt.Printf("Error reading directory: %s\n", err)
				return
			}

			for i, entry := range entries {
				isLast := i == len(entries)-1
				childPrefix := prefix + "│   "
				if isLast {
					childPrefix = prefix + "    "
				}
				printTree(filepath.Join(path, entry.Name()), childPrefix, maxNameLen, sizeMap, files, maxDepth, currentDepth+1)
			}
		}
	} else {
		// 查找文件大小
		var size int64
		for _, info := range files {
			if info.name == path {
				size = info.size
				break
			}
		}
		fmt.Printf("%sfile - %-*s - %s\n", prefix, maxNameLen, name, formatSize(size))
	}
}
